# 首先安装必要的库（如果尚未安装）
# 在命令行中运行以下命令：
# pip install scikit-learn matplotlib numpy tqdm

# 导入必要的库和模块
import numpy as np

try:
    from sklearn.datasets import load_digits
    from sklearn.model_selection import train_test_split
    from sklearn.neighbors import KNeighborsClassifier
    from sklearn.metrics import accuracy_score
except ImportError:
    print("错误：缺少必要的库。请在命令行中运行：pip install scikit-learn")
    exit(1)

try:
    import matplotlib.pyplot as plt
except ImportError:
    print("警告：matplotlib未安装，将无法显示图表")
    plt = None

try:
    from tqdm import tqdm

    has_tqdm = True
except ImportError:
    print("警告：tqdm未安装，将使用简单进度显示")
    has_tqdm = False

import pickle

# 加载数字数据集
digits = load_digits()
X = digits.data
y = digits.target

print(f"数据集形状: {X.shape}")
print(f"目标值形状: {y.shape}")
print(f"数字类别: {np.unique(y)}")

# 将数据集划分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

print(f"训练集形状: {X_train.shape}")
print(f"测试集形状: {X_test.shape}")

# 初始化变量以存储最佳准确率，相应的k值和最佳knn模型
best_accuracy = 0
best_k = 0
best_knn_model = None

# 初始化一个列表以存储每个k值的准确率
accuracy_scores = []

print("开始寻找最优K值...")
print("-" * 50)

# 尝试从1到40的k值，对于每个k值，训练knn模型，保存最佳准确率，k值和knn模型
k_range = range(1, 41)
if has_tqdm:
    k_iterator = tqdm(k_range, desc="搜索最优K值")
else:
    k_iterator = k_range
    print("进度: ", end="")

for k in k_iterator:
    # 创建KNN分类器
    knn = KNeighborsClassifier(n_neighbors=k)

    # 训练模型
    knn.fit(X_train, y_train)

    # 预测测试集
    y_pred = knn.predict(X_test)

    # 计算准确率
    accuracy = accuracy_score(y_test, y_pred)
    accuracy_scores.append(accuracy)

    # 更新最佳模型
    if accuracy > best_accuracy:
        best_accuracy = accuracy
        best_k = k
        best_knn_model = knn

    # 如果没有tqdm，显示简单进度
    if not has_tqdm:
        print(f"{k}", end=" " if k < 40 else "\n")

print("-" * 50)

# 将最佳KNN模型保存到二进制文件
with open('best_knn_model.pkl', 'wb') as file:
    pickle.dump(best_knn_model, file)

# 打印最佳准确率和相应的k值
print(f"\n🎯 最佳K值: {best_k}")
print(f"📊 最佳准确率: {best_accuracy:.4f}")
print("💾 最佳KNN模型已保存到 'best_knn_model.pkl'")

# 绘制准确率随K值变化的图表并保存为PDF
if plt is not None:
    try:
        # 设置中文字体（如果需要显示中文）
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False

        plt.figure(figsize=(10, 6))

        # 绘制折线图
        plt.plot(range(1, 41), accuracy_scores, marker='o', linestyle='-', color='b',
                 linewidth=2, markersize=4, label='Accuracy')

        # 绘制垂直线表示最佳K值
        plt.axvline(x=best_k, color='r', linestyle='--', linewidth=2,
                    label=f'Best k (k={best_k})')

        # 标记最佳点
        best_point_x = best_k
        best_point_y = best_accuracy
        plt.plot(best_point_x, best_point_y, 'ro', markersize=8, markeredgewidth=2,
                 markeredgecolor='red', markerfacecolor='yellow')

        # 添加文本标注
        plt.text(best_point_x + 0.5, best_point_y - 0.01,
                 f'k={best_point_x}\nacc={best_point_y:.4f}',
                 fontsize=10, bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))

        # 设置坐标轴标签和标题
        plt.xlabel('k value', fontsize=12)
        plt.ylabel('Accuracy', fontsize=12)
        plt.title('Accuracy of different k values', fontsize=14)

        # 设置坐标轴范围
        plt.xlim(0, 41)
        plt.ylim(min(accuracy_scores) - 0.01, max(accuracy_scores) + 0.01)

        # 添加网格和图例
        plt.grid(True, alpha=0.3)
        plt.legend()

        # 调整布局并保存为PDF
        plt.tight_layout()
        plt.savefig('accuracy_plot.pdf', format='pdf', bbox_inches='tight')
        print("📈 准确率图表已保存到 'accuracy_plot.pdf'")

        # 显示图表
        plt.show()

    except Exception as e:
        print(f"绘制图表时出错: {e}")
else:
    print("ℹ️  要查看准确率图表，请安装matplotlib: pip install matplotlib")

# 保存准确率结果到文本文件（可选）
with open('knn_results.txt', 'w', encoding='utf-8') as f:
    f.write("KNN模型寻优结果\n")
    f.write("=" * 30 + "\n")
    f.write(f"最佳K值: {best_k}\n")
    f.write(f"最佳准确率: {best_accuracy:.4f}\n")
    f.write("\n所有K值的准确率:\n")
    for k, acc in enumerate(accuracy_scores, 1):
        f.write(f"K={k:2d}: {acc:.4f}\n")

print("📄 详细结果已保存到 'knn_results.txt'")